use anyhow::{Result, anyhow, bail};
use candid::{Encode, Nat, Principal};
use canister_test::{Canister, Runtime, Wasm};
use dfn_candid::candid_one;
use ic_base_types::CanisterId;
use ic_consensus_system_test_utils::rw_message::install_nns_with_customizations_and_check_progress;
use ic_ledger_suite_orchestrator::candid::{
    AddErc20Arg, Erc20Contract, InitArg, LedgerInitArg, ManagedCanisterIds, OrchestratorArg,
    UpgradeArg,
};
use ic_nervous_system_clients::canister_status::CanisterStatusResult;
use ic_nns_constants::{GOVERNANCE_CANISTER_ID, ROOT_CANISTER_ID};
use ic_nns_test_utils::governance::{
    install_nns_canister_by_proposal, upgrade_nns_canister_with_args_by_proposal,
};
use ic_registry_subnet_type::SubnetType;
use ic_system_test_driver::driver::group::SystemTestGroup;
use ic_system_test_driver::driver::ic::{InternetComputer, Subnet};
use ic_system_test_driver::driver::test_env::TestEnv;
use ic_system_test_driver::driver::test_env_api::{
    HasPublicApiUrl, HasTopologySnapshot, IcNodeContainer, NnsCustomizations, get_dependency_path,
};
use ic_system_test_driver::systest;
use ic_system_test_driver::util::{block_on, runtime_from_url};
use ic_wasm_types::CanisterModule;
use slog::info;
use std::{future::Future, path::Path, time::Duration};

fn main() -> Result<()> {
    SystemTestGroup::new()
        .with_setup(setup_with_system_and_application_subnets)
        .add_test(systest!(ic_xc_ledger_suite_orchestrator_test))
        .execute_from_args()?;
    Ok(())
}

fn setup_with_system_and_application_subnets(env: TestEnv) {
    InternetComputer::new()
        .add_subnet(Subnet::new(SubnetType::System).add_nodes(1))
        .add_subnet(Subnet::new(SubnetType::Application).add_nodes(1))
        .setup_and_start(&env)
        .expect("Failed to setup IC under test");
    install_nns_with_customizations_and_check_progress(
        env.topology_snapshot(),
        NnsCustomizations::default(),
    );

    env.topology_snapshot()
        .subnets()
        .for_each(|subnet| subnet.await_all_nodes_healthy().unwrap());
}

fn ic_xc_ledger_suite_orchestrator_test(env: TestEnv) {
    let logger = env.logger();
    let topology_snapshot = env.topology_snapshot();

    let system_subnet_runtime = {
        let nns_subnet = topology_snapshot.root_subnet();
        let nns_node = nns_subnet.nodes().next().unwrap();
        runtime_from_url(nns_node.get_public_url(), nns_node.effective_canister_id())
    };
    let root_canister = Canister::new(&system_subnet_runtime, ROOT_CANISTER_ID);
    let governance_canister = Canister::new(&system_subnet_runtime, GOVERNANCE_CANISTER_ID);

    let application_subnet_runtime = {
        let application_subnet = topology_snapshot
            .subnets()
            .find(|s| s.subnet_type() == SubnetType::Application)
            .expect("missing application subnet");
        let application_node = application_subnet.nodes().next().unwrap();
        runtime_from_url(
            application_node.get_public_url(),
            application_node.effective_canister_id(),
        )
    };

    let ledger_orchestrator_wasm = wasm_from_path(
        "rs/ethereum/ledger-suite-orchestrator/ledger_suite_orchestrator_canister.wasm.gz",
    );
    let ledger_orchestrator = block_on(async {
        use std::str::FromStr;
        let init_args = OrchestratorArg::InitArg(InitArg {
            more_controller_ids: vec![ROOT_CANISTER_ID.get().0],
            minter_id: Some(Principal::from_str("sv3dd-oaaaa-aaaar-qacoa-cai").unwrap()),
            cycles_management: None,
        });
        let canister = install_nns_controlled_canister(
            &logger,
            &application_subnet_runtime,
            &governance_canister,
            &root_canister,
            ledger_orchestrator_wasm.clone(),
            Encode!(&init_args).unwrap(),
        )
        .await;
        LedgerOrchestratorCanister { canister }
    });
    info!(
        &logger,
        "Installed ledger orchestrator canister at {}",
        ledger_orchestrator.as_ref().canister_id()
    );

    block_on(async {
        upgrade_ledger_suite_orchestrator_by_nns_proposal(
            &logger,
            &governance_canister,
            &root_canister,
            ledger_orchestrator_wasm.clone(),
            &ledger_orchestrator,
            OrchestratorArg::UpgradeArg(UpgradeArg {
                git_commit_hash: Some("6a8e5fca2c6b4e12966638c444e994e204b42989".to_string()),
                ..Default::default()
            }),
        )
        .await
    });
    info!(
        &logger,
        "Registered embedded wasms in the ledger suite orchestrator {}",
        ledger_orchestrator.as_ref().canister_id()
    );

    let managed_canister_ids =
        block_on(async { ledger_orchestrator.call_canister_ids(usdc_contract()).await });
    assert_eq!(managed_canister_ids, None);

    let usdc_ledger_suite = block_on(async {
        add_erc_20_by_nns_proposal(
            &logger,
            &governance_canister,
            &root_canister,
            ledger_orchestrator_wasm,
            &ledger_orchestrator,
            AddErc20Arg {
                contract: usdc_contract(),
                ledger_init_arg: usdc_ledger_init_arg(),
            },
        )
        .await
    });

    block_on(async {
        let token_name: String = try_async("getting token_name", &logger, || {
            usdc_ledger_suite
                .ledger
                .query_("icrc1_name", candid_one, ())
        })
        .await;
        assert_eq!(token_name, "USD Coin".to_string());
    });
    info!(
        &logger,
        "USDC ledger {} is up and running",
        usdc_ledger_suite.ledger.canister_id(),
    );

    block_on(async {
        let index_status: ic_icrc1_index_ng::Status = try_async("getting status", &logger, || {
            usdc_ledger_suite.index.query_("status", candid_one, ())
        })
        .await;
        assert_eq!(
            index_status,
            ic_icrc1_index_ng::Status {
                num_blocks_synced: Nat::from(0_u8)
            }
        );
    });
    info!(
        &logger,
        "USDC index {} is up and running",
        usdc_ledger_suite.index.canister_id(),
    );

    assert!(usdc_ledger_suite.archives.is_empty());
}

async fn install_nns_controlled_canister<'a>(
    logger: &slog::Logger,
    application_subnet_runtime: &'a Runtime,
    governance_canister: &Canister<'_>,
    root_canister: &Canister<'_>,
    canister_wasm: CanisterModule,
    canister_init_payload: Vec<u8>,
) -> Canister<'a> {
    use ic_nervous_system_clients::canister_status::CanisterStatusType;

    let canister = application_subnet_runtime
        .create_canister(Some(u128::MAX))
        .await
        .expect("failed to create canister");
    info!(
        logger,
        "Created empty canister at {}",
        canister.canister_id()
    );

    canister
        .set_controller(ROOT_CANISTER_ID.get())
        .await
        .expect("failed to modify canister controller");
    info!(
        logger,
        "Change controller of {} to root {}",
        canister.canister_id(),
        ROOT_CANISTER_ID
    );

    let new_module_hash = canister_wasm.module_hash();
    let wasm = Wasm::from_bytes(canister_wasm.as_slice());
    install_nns_canister_by_proposal(
        &canister,
        governance_canister,
        root_canister,
        wasm,
        Some(canister_init_payload),
    )
    .await;
    info!(
        logger,
        "Installed WASM to {} via NNS proposal",
        canister.canister_id()
    );

    status_of_nns_controlled_canister_satisfy(logger, root_canister, &canister, |status| {
        status.status == CanisterStatusType::Running
            && status.module_hash.as_deref() == Some(new_module_hash.as_slice())
    })
    .await;

    info!(logger, "Canister {} is ready!", canister.canister_id());

    canister
}

async fn upgrade_ledger_suite_orchestrator_by_nns_proposal(
    logger: &slog::Logger,
    governance_canister: &Canister<'_>,
    root_canister: &Canister<'_>,
    canister_wasm: CanisterModule,
    orchestrator: &LedgerOrchestratorCanister<'_>,
    upgrade_arg: OrchestratorArg,
) {
    use ic_nervous_system_clients::canister_status::CanisterStatusType;

    upgrade_nns_canister_with_args_by_proposal(
        orchestrator.as_ref(),
        governance_canister,
        root_canister,
        true,
        Wasm::from_bytes(canister_wasm.as_slice()),
        Encode!(&upgrade_arg).unwrap(),
    )
    .await;
    info!(
        logger,
        "Upgrade ledger suite orchestrator {:?} via NNS proposal", upgrade_arg
    );

    status_of_nns_controlled_canister_satisfy(
        logger,
        root_canister,
        orchestrator.as_ref(),
        |status| status.status == CanisterStatusType::Running,
    )
    .await;

    info!(
        logger,
        "Upgrade finished. Ledger orchestrator is back running"
    );
}

async fn add_erc_20_by_nns_proposal<'a>(
    logger: &slog::Logger,
    governance_canister: &Canister<'_>,
    root_canister: &Canister<'_>,
    canister_wasm: CanisterModule,
    orchestrator: &LedgerOrchestratorCanister<'a>,
    erc20_token: AddErc20Arg,
) -> ManagedCanisters<'a> {
    let erc20_contract = erc20_token.contract.clone();
    upgrade_ledger_suite_orchestrator_by_nns_proposal(
        logger,
        governance_canister,
        root_canister,
        canister_wasm,
        orchestrator,
        OrchestratorArg::AddErc20Arg(erc20_token),
    )
    .await;

    let created_canister_ids = ic_system_test_driver::retry_with_msg_async!(
        "checking if all canisters are created",
        logger,
        Duration::from_secs(100),
        Duration::from_secs(1),
        || async {
            let managed_canister_ids = orchestrator.call_canister_ids(erc20_contract.clone()).await;
            match managed_canister_ids {
                None => bail!("No managed canister IDs yet"),
                Some(x) if x.ledger.is_some() && x.index.is_some() => Ok(x),
                _ => bail!(
                    "Not all canisters were created yet: {:?}",
                    managed_canister_ids
                ),
            }
        }
    )
    .await
    .unwrap_or_else(|e| panic!("Canisters for contract {erc20_contract:?} were not created: {e}"));
    info!(
        &logger,
        "Created canister IDs: {} for contract {:?}", created_canister_ids, erc20_contract
    );

    ManagedCanisters::from(orchestrator.as_ref().runtime(), created_canister_ids)
}

fn wasm_from_path<P: AsRef<Path>>(path: P) -> CanisterModule {
    CanisterModule::new(Wasm::from_file(get_dependency_path(path)).bytes())
}

fn usdc_contract() -> Erc20Contract {
    Erc20Contract {
        chain_id: Nat::from(1_u8),
        address: "0xa0b86991c6218b36c1d19d4a2e9eb0ce3606eb48".to_string(),
    }
}

fn usdc_ledger_init_arg() -> LedgerInitArg {
    const CKETH_TOKEN_LOGO: &str = "";

    LedgerInitArg {
        transfer_fee: 2_000_000_000_000_u64.into(),
        decimals: 6,
        token_name: "USD Coin".to_string(),
        token_symbol: "USDC".to_string(),
        token_logo: CKETH_TOKEN_LOGO.to_string(),
    }
}

async fn status_of_nns_controlled_canister_satisfy<P: Fn(&CanisterStatusResult) -> bool>(
    logger: &slog::Logger,
    root_canister: &Canister<'_>,
    target_canister: &Canister<'_>,
    predicate: P,
) {
    use dfn_candid::candid;

    ic_system_test_driver::retry_with_msg_async!(
        format!(
            "calling canister_status of {} to check if {} satisfies the predicate",
            root_canister.canister_id(),
            target_canister.canister_id()
        ),
        logger,
        Duration::from_secs(60),
        Duration::from_secs(1),
        || async {
            let status: CanisterStatusResult = root_canister
                .update_("canister_status", candid, (target_canister.as_record(),))
                .await
                .map_err(|e| anyhow!(e))?;
            info!(
                logger,
                "Canister status of {}: {:?}",
                target_canister.canister_id(),
                status
            );
            if predicate(&status) {
                Ok(())
            } else {
                bail!(
                    "Status of {} did not satisfy predicate",
                    target_canister.canister_id()
                )
            }
        }
    )
    .await
    .unwrap_or_else(|e| {
        panic!(
            "Canister status of {} did not satisfy predicate: {}",
            target_canister.canister_id(),
            e
        )
    });
}

struct LedgerOrchestratorCanister<'a> {
    canister: Canister<'a>,
}

impl LedgerOrchestratorCanister<'_> {
    async fn call_canister_ids(&self, contract: Erc20Contract) -> Option<ManagedCanisterIds> {
        self.canister
            .query_("canister_ids", dfn_candid::candid, (contract,))
            .await
            .expect("Error while calling canister_ids endpoint")
    }
}

impl<'a> AsRef<Canister<'a>> for LedgerOrchestratorCanister<'a> {
    fn as_ref(&self) -> &Canister<'a> {
        &self.canister
    }
}

struct ManagedCanisters<'a> {
    ledger: Canister<'a>,
    index: Canister<'a>,
    archives: Vec<Canister<'a>>,
}

impl<'a> ManagedCanisters<'a> {
    pub fn from(runtime: &'a Runtime, canister_ids: ManagedCanisterIds) -> Self {
        let to_canister_id =
            |canister_id: Principal| CanisterId::unchecked_from_principal(canister_id.into());

        Self {
            ledger: Canister::new(runtime, to_canister_id(canister_ids.ledger.unwrap())),
            index: Canister::new(runtime, to_canister_id(canister_ids.index.unwrap())),
            archives: canister_ids
                .archives
                .into_iter()
                .map(|archive_id| Canister::new(runtime, to_canister_id(archive_id)))
                .collect(),
        }
    }
}

async fn try_async<S: AsRef<str>, F, Fut, R>(msg: S, logger: &slog::Logger, f: F) -> R
where
    Fut: Future<Output = Result<R, String>>,
    F: Fn() -> Fut,
{
    ic_system_test_driver::retry_with_msg_async!(
        msg.as_ref(),
        logger,
        Duration::from_secs(100),
        Duration::from_secs(1),
        || async { f().await.map_err(|e| anyhow!(e)) }
    )
    .await
    .expect("failed despite retries")
}
